import gevent
import time
import json
import numpy as np
import pandas as pd
import os, sys
sys.path.append('..')
from contract_market.server.server import MainProcess  # NOQA
from contract_market.utils.api import XTAPI
from contract_market.initialize.config import InitConfig  # NOQA
from contract_market.common.constant import keys  # NOQA
from contract_market.utils.redisdb import Redis
from contract_market.utils.decorator import loop, exception
from contract_market.main import GracefulKiller
from contract_market.server.server import Argv


timeout = 0.1

class Restruct(object):

    def __init__(self) -> None:
        self.columns = ['state', 'price', 'origQty', 'executedQty', 'orderId', 'orderSide']
        self.maker = pd.DataFrame()
        self.balance = None
        self.index_price = None
        self.mark_price = None
        self.position = None

    def init_maker(self, xt_api, symbol):
        if maker.empty or not int(time.time())%5:
            data = xt_api.open_orders_api(symbol)
            if not data.empty:
                maker = self.oorder_to_orderbook(data)
            else:
                maker = pd.DataFrame()
            gevent.sleep(0.1)
        self.maker = maker
        return True

    def oorder_to_orderbook(self, data):
        # data['orderId'] = data.orderId
        data = data[self.columns]
        data = data.astype({'price': float, 'origQty': float, 'executedQty': float, 'orderId': int})
        self.maker = data
        return self.maker

    def orderbook(self, xt_api, symbol, pub_key, res):
        '''
        -------------------------------------------
        |STATE             |    ACTION            |
        -------------------------------------------
        |NEW               |    insert or delete  |
        |                  |                      |
        |CANCELED          |    remove order      |
        |EXPIRED           |    remove order      |
        |REJECTED          |    remove order      |
        |FILLED            |    remove order      |
        |                  |                      |
        |PARTIALLY_FILLED  |    change quantity   |
        |PARTIALLY_CANCELED|    change quantity   |
        -------------------------------------------
        '''
        self.init_maker(xt_api, symbol)

        if res is None:
            return False
        # check public key
        if pub_key not in res.get('channel').decode('utf-8'):
            return False
        data = json.loads(res.get('data'))
        # check symbol
        if data.get('symbol') != symbol:
            return False
        # 
        data = pd.DataFrame([data]).astype({
            'price': float, 'origQty': float, 'executedQty': float, 'orderId': int}).squeeze()
        state = data.state
        create_state = ['NEW']
        remove_state = ['CANCELED', 'EXPIRED', 'FILLED']
        change_state = ['PARTIALLY_FILLED', 'PARTIALLY_CANCELED']
        unknow_state = ['REJECTED']
        if state in create_state:
            self.maker = self.maker.append(data)
        # remove
        elif state in remove_state and not self.maker.empty:
            self.maker = self.maker[self.maker.orderId!=data.orderId]
        # unknow
        elif state in unknow_state and not self.maker.empty:
            print('NotImplementedError:', state)
        # change quantity
        elif state in change_state and not self.maker.empty:
            change = self.maker[self.maker.orderId==data.orderId]
            if not change.empty:
                self.maker.loc[change.index, 'origQty'] = change.origQty-data.executedQty
        else:
            print('NotImplementedError: xxx')
        # update
        if not self.maker.empty:
            self.maker = self.maker[self.maker.origQty!=0]
            self.maker = self.maker.reset_index(drop=True)
            print('maker', len(self.maker), len(self.maker[self.maker.orderSide=='BUY']), \
                len(self.maker[self.maker.orderSide=='SELL']))
        maker = maker.drop_duplicates()
        return self.maker


def subscribe(engine, redis, symbol, pub_key, transcation_type, verbose=True):
    
    @exception
    def _replace_orderbook_job(restruct, xt_api, queue, key):
        maker = None
        if restruct.maker.empty or not int(time.time())%30:
            data = xt_api.open_orders_api(symbol, page=10)
            if not data.empty:
                maker = restruct.oorder_to_orderbook(data)
            else:
                maker = pd.DataFrame()
            gevent.sleep(0.1)
        # save maker
        if maker is not None and not isinstance(maker, bool):
            print('replace orderbook')
            maker_series = maker.to_json(orient='records')
            redis.standalone_operator('set', key, maker_series)
        return True
    
    def replace_orderbook_job(symbol, pub_key):
        print('subscribe orderbook:', symbol)
        key = f'{symbol}:orderbook'
        restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        queue = engine.on_user_trade(pub_key)
        while True:
            _replace_orderbook_job(restruct, xt_api, queue, key)
            gevent.sleep(1)

    @exception
    def _orderbook_job(restruct, xt_api, queue, key):
        gevent.sleep(0.1)
        # update maker 
        maker = None
        if restruct.maker.empty:
            data = xt_api.open_orders_api(symbol)
            if not data.empty:
                maker = restruct.oorder_to_orderbook(data)
            else:
                maker = pd.DataFrame()
            gevent.sleep(0.1)
        else:
            try:
                res = queue.get_notwait()
                if verbose and res:
                    print('orderbook', res)
            except Exception as ex:
                res = None
                print(ex)
            maker = restruct.orderbook(xt_api, symbol, pub_key, res)
        # save maker
        if maker is not None and not isinstance(maker, bool):
            maker_series = maker.to_json(orient='records')
            redis.standalone_operator('set', key, maker_series)
        return True

    def orderbook_job(symbol, pub_key):
        print('subscribe orderbook:', symbol)
        key = f'{symbol}:orderbook'
        restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        queue = engine.on_user_trade(pub_key)
        while True:
            _orderbook_job(restruct, xt_api, queue, key)
    
    @exception
    def _open_orders_asyn_job(symbol, xt_api, key, pages):
        '''
        return 
            price, quantity: [float], [float]
        '''
        @exception
        def job(symbol, page):
            gevent.sleep(timeout)
            order = xt_api.open_orders_single_api(symbol, page=page)
            redis.standalone_operator('set', key+str(page), json.dumps(order))
        
        jobs = [gevent.spawn(job, **{'symbol': symbol, 'page': page}) for page in range(1, pages+1)]
        gevent.joinall(jobs, timeout=3)
        return True

    def open_orders_asyn_job(symbol, pub_key, pages=5):
        print('subscribe open orders:', symbol)
        key = f'{symbol}:orders'
        # restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        # queue = engine.on_user_balance(pub_key)
        while True:
            _open_orders_asyn_job(symbol, xt_api, key, pages)

    @exception
    def _balance_job(restruct, xt_api, queue, key):
        gevent.sleep(5)
        balance = ''
        if not restruct.balance or not int(time.time())%1:
            unit = symbol.split('_')[1] if transcation_type == 'um' else symbol.split('_')[0]
            balance = xt_api.balance_api(unit=unit)
            restruct.balance = balance
            gevent.sleep(1)
        # else:
        #     try:
        #         res = queue.get_notwait()
        #         if verbose and res:
        #             print('balance', res)
        #     except Exception as ex:
        #         res = None
        #     if res:
        #         data = json.loads(res.get('data'))
        #         balance = float(data.get('availableBalance'))
        if balance:
            redis.standalone_operator('set', key, str(balance))
        return True

    def balance_job(symbol, pub_key):
        print('subscribe balance:', symbol)
        key = f'{pub_key}:balance'
        restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        queue = engine.on_user_balance(pub_key)
        while True:
            _balance_job(restruct, xt_api, queue, key)
    
    @exception
    def _index_price_job(restruct, xt_api, key):
        gevent.sleep(0.1)
        index_price = ''
        if not restruct.index_price or not int(time.time())%1:
            index_price = xt_api.index_price_api(symbol)
            print('api index price', index_price)
            restruct.index_price = index_price
            gevent.sleep(0.1)
        # else:
        #     try:
        #         res = queue.get_notwait()
        #         if verbose and res:
        #             print('index price', res)
        #     except Exception as ex:
        #         res = None
        #     if res:
        #         data = json.loads(res.get('data'))
        #         index_price = data.get('data').get('p')
        if index_price:
            redis.standalone_operator('set', key, str(index_price))
        return True

    def index_price_job(symbol, pub_key):
        print('subscribe index price:', symbol)
        key = f'{symbol}:index_price'
        restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        # queue = engine.on_index_price(symbol)
        while True:
            _index_price_job(restruct, xt_api, key)
    
    @exception
    def _mark_price_job(restruct, xt_api, key):
        gevent.sleep(0.1)
        mark_price = ''
        if not restruct.mark_price or not int(time.time())%1:
            mark_price = xt_api.mark_price_api(symbol)
            print('api mark price', mark_price)
            restruct.mark_price = mark_price
            gevent.sleep(0.1)
        if mark_price:
            redis.standalone_operator('set', key, str(mark_price))
        return True

    def mark_price_job(symbol, pub_key):
        print('subscribe mark price:', symbol)
        key = f'{symbol}:mark_price'
        restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        while True:
            _mark_price_job(restruct, xt_api, key)

    @exception
    def _user_position(restruct, xt_api, queue, key):
        gevent.sleep(0.1)
        pos = ''
        if not restruct.position or not int(time.time())%1:
            pos = xt_api.position(symbol, 'api')
            restruct.position = pos
            gevent.sleep(0.1)
        if pos:
            redis.standalone_operator('set', key, str(pos))
        return True
    
    def user_position(symbol, pub_key):
        print('subscribe index price:', symbol)
        key = f'{symbol}:user_position'
        restruct = Restruct()
        xt_api = XTAPI()
        xt_api.init_client(symbol, transcation_type)
        # queue = engine.on_index_price(symbol)
        queue = None
        while True:
            _user_position(restruct, xt_api, queue, key)

    kwargs = {'symbol': symbol, 'pub_key': pub_key}
    # orderbook_job
    jobs = [gevent.spawn(job, **kwargs) for job in [index_price_job, mark_price_job, balance_job, user_position, open_orders_asyn_job]]
    gevent.joinall(jobs)
    
    # user_position(symbol, pub_key)

    # import threading
    # jobs = [index_price_job, balance_job, orderbook_job, user_position]
    # ths = [threading.Thread(target=job, kwargs=kwargs, daemon=True) for job in jobs]
    # for t in ths:
    #     t.start()
    # for t in ths:
    #     t.join()
    return True


def run_subscribe():
    '''
    TODO: 相同的timestamp调度会出问题
    '''
    # load engine
    killer = GracefulKiller()
    mp = MainProcess(killer)
    event, engine = mp.init_engine()
    redis = Redis()
    # load configure
    cfg = InitConfig("image", "pair.yaml", argv=Argv)
    print("[cfg]", cfg.resolve(), Argv.process_name)
    configure = cfg.resolve()
    
    # run
    from multiprocessing import Process

    process_list = []
    for symbol, value in configure.items():
        p = Process(target=subscribe, args=(engine, redis, symbol, keys.get(symbol)[0], value.get('transcation_type')))
        p.start()
        process_list.append(p)
    for p in process_list:
        p.join()
